from tool_yolo_other import CT_detection

model_infer = CT_detection()
import cv2



def yolo_get_point_img(im0):
    #获取数据尺寸
    datas_shape = im0.shape
    #数据前处理 转换成模型需要的尺寸
    img, im, data_info = model_infer.preprocessing_img(im0)
    #模型推理
    pred = model_infer.model(im, augment=False, visualize=False)
    #数据后处理
    pred = model_infer.non_max_suppression(pred, 0.25, 0.25, None, False, max_det=1000)
    #结果可视化
    result= model_infer.postprocessing_img(pred, img, im, datas_shape)



    return result

if __name__ =="__main__":
    import os
    now_path = os.getcwd()
    input_path = os.path.join(now_path,'input_dir')
    output_path = os.path.join(now_path,'output_dir')
    if os.path.exists(output_path) is False:
        os.makedirs(output_path)
    files = os.listdir(input_path)
    for file in files:
        file_path = os.path.join(input_path,file)
        save_path = os.path.join(output_path,file)
        result = yolo_get_point_img(cv2.imread(file_path))


        cv2.imwrite(save_path,result)



